// Copyright (c) Lawrence Livermore National Security, LLC and other VisIt
// Project developers.  See the top-level LICENSE file for dates and other
// details.  No copyright assignment is required to contribute to VisIt.

// ****************************************************************************
//  avtHTGWriter.C
// ****************************************************************************

#include <avtHTGWriter.h>

#include <vector>

#include <vtkCellData.h>
#include <vtkDataArray.h>
#include <vtkDataSetWriter.h>
#include <vtkRectilinearGrid.h>

#include <avtDatabaseMetaData.h>

#include <avtHTGOptions.h>
#include <DBOptionsAttributes.h>
#include <DebugStream.h>

#include <ImproperUseException.h>

using     std::string;
using     std::vector;

static float HTGCreate(const float *varIn, float *varOut, int *mask,
    int nLevels, int nx, int blankValue, int level, int *offsets,
    int iStart, int jStart, int kStart);
static void HTGWriteFile(const string &stem, const double *bounds,
    int nLevels, int nVertices, int nDescriptor, int descriptorMin,
    int descriptorMax, const int *descriptor, int nbVerticesByLevelMax,
    const int *nbVerticesByLevel, int nMask, int maskMin, int maskMax,
    const int *mask, double varMin, double varMax, const float *var);
static void HTGWrite(const std::string &path, float blankValue, int nx,
    const double *bounds, const float *value);

// ****************************************************************************
//  Method: avtHTGWriter constructor
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Mon May 1 15:28:30 PST 2023
//
// ****************************************************************************

avtHTGWriter::avtHTGWriter(const DBOptionsAttributes *atts)
{
    blankValue = atts->GetDouble(HTGDBOptions::HTG_BLANK_VALUE);
}

// ****************************************************************************
//  Method: avtHTGWriter::OpenFile
//
//  Purpose:
//      Does no actual work.  Just records the stem name for the files.
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Mon May 1 15:28:30 PST 2023
//
// ****************************************************************************

void
avtHTGWriter::OpenFile(const string &stemname, int numblocks)
{
    stem = stemname;

    if (numblocks > 1)
    {
        EXCEPTION1(ImproperUseException,
            "The HTG writer can only handle single block files.");
    }
}


// ****************************************************************************
//  Method: avtHTGWriter::WriteHeaders
//
//  Purpose:
//      Writes out a VisIt file to tie the HTG files together.
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Mon May 1 15:28:30 PST 2023
//
// ****************************************************************************

void
avtHTGWriter::WriteHeaders(const avtDatabaseMetaData *md,
                           const vector<string> &scalars,
                           const vector<string> &vectors,
                           const vector<string> &materials)
{
    // 
    // Since we only support writing out a single block there is no
    // advantage to writing out a header ahead of time, so just do it
    // all in WriteChunk.
    // 
}


// ****************************************************************************
//  Method: avtHTGWriter::WriteChunk
//
//  Purpose:
//      This writes out one chunk of an avtDataset.
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Mon May 1 15:28:30 PST 2023
//
// ****************************************************************************

void
avtHTGWriter::WriteChunk(vtkDataSet *ds, int chunk)
{
    int   i;
    if (ds->GetDataObjectType() != VTK_RECTILINEAR_GRID)
    {
        EXCEPTION1(ImproperUseException,
            "The HTG writer can only handle rectilinear grids.");
    }
    vtkRectilinearGrid *rgrid = (vtkRectilinearGrid *) ds;

    int dims[3];
    rgrid->GetDimensions(dims);

    //
    // Check that the dimensions are all equal.
    //
    int nx, ny, nz;
    nx = dims[0];
    ny = dims[1];
    nz = dims[2];
    if (nx != ny && ny != nz)
    {
        EXCEPTION1(ImproperUseException,
            "The HTG writer requires the dimensions to be equal.");
    }

    //
    // Check that the dimension is a power of 2.
    //
    nx = nx - 1;
    if (nx == 0 || ((nx & (nx - 1)) != 0))
    {
        EXCEPTION1(ImproperUseException,
            "The HTG writer requires the grid dimension to be a power of 2.");
    }

    //
    // Check that the spacing is uniform in each direction.
    //
    vtkDataArray *coords[3];
    coords[0] = rgrid->GetXCoordinates();
    coords[1] = rgrid->GetYCoordinates();
    coords[2] = rgrid->GetZCoordinates();
    for (int i = 0; i < 3; i++)
    {
        vtkDataArray *coord = coords[i];
        double dx = (coord->GetTuple1(dims[0]-1) - coord->GetTuple1(0)) /
            double(dims[0]-1);
        double fuzz = dx / 1.0e-6;
        for (int j = 0; j < nx; j++)
        {
            double curDx = coord->GetTuple1(j+1) - coord->GetTuple1(j);
            if (curDx - dx > fuzz)
            {
                EXCEPTION1(ImproperUseException,
                    "The HTG writer requires constant spacing in each direction.");
            }
        }
    }

    double bounds[6];
    rgrid->GetBounds(bounds);

    //
    // Write out the cell arrays.
    //
    int nVars = rgrid->GetCellData()->GetNumberOfArrays();
    int nWritten = 0;
    for (int i = 0; i < nVars; i++)
    { 
        vtkDataArray *arr = rgrid->GetCellData()->GetArray(i);
        if (arr->GetNumberOfComponents() == 1)
        {
            float *values = NULL;
            bool deleteValues = false;
            if (arr->GetDataType() == VTK_FLOAT)
                values = (float *) arr->GetVoidPointer(0);
            else
            {
                int nvals = arr->GetNumberOfTuples();
                values = new float[nvals];
                for (int j = 0; j < nvals; j++)
                {
                    values[j] = (float) arr->GetTuple1(j);
                }
                deleteValues = true;
            }

            //
            // HTGWrite adds the ".htg" extension to the filename.
            //
            std::string path = stem + "." + arr->GetName();
            try
            {
                HTGWrite(path.c_str(), blankValue, nx, bounds, values);
            }
            catch (ImproperUseException)
            {
                debug1 << "The variable " << arr->GetName()
                       << " only had blank values." << endl;
            }

            if (deleteValues)
                delete [] values;

            nWritten++;
        }
    }
    if (nWritten < 1)
    {
        EXCEPTION1(ImproperUseException,
            "The HTG writer didn't find any variables to write.");
    }
}


// ****************************************************************************
//  Method: avtHTGWriter::CloseFile
//
//  Purpose:
//      Closes the file.  This does nothing in this case.
//
//  Programmer: brugger -- generated by xml2avt
//  Creation:   Mon May 1 15:28:30 PST 2023
//
// ****************************************************************************

void
avtHTGWriter::CloseFile(void)
{
    // Just needed to meet interface.
}

// ****************************************************************************
//  Method: HTGCreate
//
//  Purpose:
//      Creates a HyperTreeGrid.
//
//  Programmer: Eric Brugger
//  Creation:   Mon May 1 15:28:30 PST 2023
//
// ****************************************************************************

float HTGCreate(const float *varIn, float *varOut, int *mask, int nLevels,
    int nx, int blankValue, int level, int *offsets,
    int iStart, int jStart, int kStart)
{
    if (level < nLevels - 1)
    {
        //
        // Recurse over the 8 sub-blocks.
        //
        int nxSubBox = 1 << (nLevels - level - 1);
        int offset = offsets[level];
        offsets[level] = offsets[level] + 8;

        varOut[offset]   = HTGCreate(varIn, varOut, mask, nLevels, nx,
            blankValue, level+1, offsets,
            iStart, jStart, kStart);
        varOut[offset+1] = HTGCreate(varIn, varOut, mask, nLevels, nx,
            blankValue, level+1, offsets,
            iStart + nxSubBox, jStart, kStart);
        varOut[offset+2] = HTGCreate(varIn, varOut, mask, nLevels, nx,
            blankValue, level+1, offsets,
            iStart, jStart + nxSubBox, kStart);
        varOut[offset+3] = HTGCreate(varIn, varOut, mask, nLevels, nx,
            blankValue, level+1, offsets,
            iStart + nxSubBox, jStart + nxSubBox, kStart);
        varOut[offset+4] = HTGCreate(varIn, varOut, mask, nLevels, nx,
            blankValue, level+1, offsets,
            iStart, jStart, kStart + nxSubBox);
        varOut[offset+5] = HTGCreate(varIn, varOut, mask, nLevels, nx,
            blankValue, level+1, offsets,
            iStart + nxSubBox, jStart, kStart + nxSubBox);
        varOut[offset+6] = HTGCreate(varIn, varOut, mask, nLevels, nx,
            blankValue, level+1, offsets,
            iStart, jStart + nxSubBox, kStart + nxSubBox);
        varOut[offset+7] = HTGCreate(varIn, varOut, mask, nLevels, nx,
            blankValue, level+1, offsets,
            iStart + nxSubBox, jStart + nxSubBox, kStart + nxSubBox);
        mask[offset]   = varOut[offset]   == blankValue ? 1 : 0;
        mask[offset+1] = varOut[offset+1] == blankValue ? 1 : 0;
        mask[offset+2] = varOut[offset+2] == blankValue ? 1 : 0;
        mask[offset+3] = varOut[offset+3] == blankValue ? 1 : 0;
        mask[offset+4] = varOut[offset+4] == blankValue ? 1 : 0;
        mask[offset+5] = varOut[offset+5] == blankValue ? 1 : 0;
        mask[offset+6] = varOut[offset+6] == blankValue ? 1 : 0;
        mask[offset+7] = varOut[offset+7] == blankValue ? 1 : 0;

        //
        // Calculate and return the average.
        //
        float ave = 0.;
        int nVal = 0;
        for (int l = 0; l < 8; l++)
        {
            if (varOut[offset+l] != blankValue)
            {
                nVal++;
                ave += varOut[offset+l];
            }
        }
        if (nVal)
            ave /= float(nVal);
        else
            ave = blankValue;
        return ave;
    }
    else
    {
        //
        // Base case. Reorder and store the values as the finest
        // resolution.
        //
        int offset = offsets[level];
        offsets[level] = offsets[level] + 8;

        int index  = (kStart)   * nx * nx + (jStart)   * nx + (iStart);
        int index2 = (kStart)   * nx * nx + (jStart)   * nx + (iStart+1);
        int index3 = (kStart)   * nx * nx + (jStart+1) * nx + (iStart);
        int index4 = (kStart)   * nx * nx + (jStart+1) * nx + (iStart+1);
        int index5 = (kStart+1) * nx * nx + (jStart)   * nx + (iStart);
        int index6 = (kStart+1) * nx * nx + (jStart)   * nx + (iStart+1);
        int index7 = (kStart+1) * nx * nx + (jStart+1) * nx + (iStart);
        int index8 = (kStart+1) * nx * nx + (jStart+1) * nx + (iStart+1);
        varOut[offset]   = varIn[index];
        varOut[offset+1] = varIn[index2];
        varOut[offset+2] = varIn[index3];
        varOut[offset+3] = varIn[index4];
        varOut[offset+4] = varIn[index5];
        varOut[offset+5] = varIn[index6];
        varOut[offset+6] = varIn[index7];
        varOut[offset+7] = varIn[index8];
        mask[offset]   = varOut[offset]   == blankValue ? 1 : 0;
        mask[offset+1] = varOut[offset+1] == blankValue ? 1 : 0;
        mask[offset+2] = varOut[offset+2] == blankValue ? 1 : 0;
        mask[offset+3] = varOut[offset+3] == blankValue ? 1 : 0;
        mask[offset+4] = varOut[offset+4] == blankValue ? 1 : 0;
        mask[offset+5] = varOut[offset+5] == blankValue ? 1 : 0;
        mask[offset+6] = varOut[offset+6] == blankValue ? 1 : 0;
        mask[offset+7] = varOut[offset+7] == blankValue ? 1 : 0;

        //
        // Calculate and return the average.
        //
        float ave = 0.;
        int nVal = 0;
        for (int l = 0; l < 8; l++)
        {
            if (varOut[offset+l] != blankValue)
            {
                nVal++;
                ave += varOut[offset+l];
            }
        }
        if (nVal)
            ave /= float(nVal);
        else
            ave = blankValue;
        return ave;
    }
}

// ****************************************************************************
//  Method: HTGWriteFile
//
//  Purpose:
//      Writes a HyperTreeGrid file.
//
//  Programmer: Eric Brugger
//  Creation:   Mon May 1 15:28:30 PST 2023
//
// ****************************************************************************

void HTGWriteFile(const string &stem, const double *bounds, int nLevels,
    int nVertices, int nDescriptor, int descriptorMin, int descriptorMax,
    const int *descriptor, int nbVerticesByLevelMax,
    const int *nbVerticesByLevel, int nMask, int maskMin, int maskMax,
    const int *mask, double varMin, double varMax, const float *var)
{
    //
    // Write out the HTG VTK file. It is in ASCII format, which is the
    // least efficient, but it's the simplest and was great for developing
    // the algorithm. This should probably be improved at some point.
    //
    std::string filename = stem + ".htg";
    ofstream *ofile = new ofstream(filename.c_str());

    *ofile << "<VTKFile type=\"HyperTreeGrid\" version=\"1.0\" byte_order=\"LittleEndian\" header_type=\"UInt32\">" << endl;
    *ofile << "  <HyperTreeGrid BranchFactor=\"2\" TransposedRootIndexing=\"0\" Dimensions=\"2 2 2\">" << endl;
    *ofile << "    <Grid>" << endl;
    *ofile << "      <DataArray type=\"Float64\" Name=\"XCoordinates\" NumberOfTuples=\"2\" format=\"ascii\" RangeMin=\"" << bounds[0] << "\" RangeMax=\"" << bounds[1] << "\">" << endl;
    *ofile << "        " << bounds[0] << " " << bounds[1] << endl;
    *ofile << "      </DataArray>" << endl;
    *ofile << "      <DataArray type=\"Float64\" Name=\"YCoordinates\" NumberOfTuples=\"2\" format=\"ascii\" RangeMin=\"" << bounds[2] << "\" RangeMax=\"" << bounds[3] << "\">" << endl;
    *ofile << "        " << bounds[2] << " " << bounds[3] << endl;
    *ofile << "      </DataArray>" << endl;
    *ofile << "      <DataArray type=\"Float64\" Name=\"ZCoordinates\" NumberOfTuples=\"2\" format=\"ascii\" RangeMin=\"" << bounds[4] << "\" RangeMax=\"" << bounds[5] << "\">" << endl;
    *ofile << "        " << bounds[4] << " " << bounds[5] << endl;
    *ofile << "      </DataArray>" << endl;
    *ofile << "    </Grid>" << endl;
    *ofile << "    <Trees>" << endl;
    *ofile << "      <Tree Index=\"0\" NumberOfLevels=\"" << nLevels << "\" NumberOfVertices=\"" << nVertices << "\">" << endl;
    *ofile << "        <DataArray type=\"Bit\" Name=\"Descriptor\" NumberOfTuples=\"" << nDescriptor << "\" format=\"ascii\" RangeMin=\"" << descriptorMin << "\" RangeMax=\"" << descriptorMax << "\">" << endl;
    for (int i = 0; i < nDescriptor; i += 6)
    {
        *ofile << "          ";
        int jmax = (i + 6 < nDescriptor) ? i + 6 : nDescriptor;
        for (int j = i; j < jmax - 1; j++)
            *ofile << descriptor[j] << " ";
        *ofile << descriptor[jmax-1] << endl;;
    }
    *ofile << "        </DataArray>" << endl;
    *ofile << "        <DataArray type=\"Int64\" Name=\"NbVerticesByLevel\" NumberOfTuples=\"" << nLevels << "\" format=\"ascii\" RangeMin=\"1\" RangeMax=\"" << nbVerticesByLevelMax << "\">" << endl;
    *ofile << "          ";
    for (int i = 0; i < nLevels - 1; i++)
        *ofile << nbVerticesByLevel[i] << " ";
    *ofile << nbVerticesByLevel[nLevels-1] << endl;;
    *ofile << "        </DataArray>" << endl;
    *ofile << "        <DataArray type=\"Bit\" Name=\"Mask\" NumberOfTuples=\"" << nMask << "\" format=\"ascii\" RangeMin=\"" << maskMin << "\" RangeMax=\"" << maskMax << "\">" << endl;
    for (int i = 0; i < nMask; i += 6)
    {
        *ofile << "          ";
        int jmax = (i + 6 < nMask) ? i + 6 : nMask;
        for (int j = i; j < jmax - 1; j++)
            *ofile << mask[j] << " ";
        *ofile << mask[jmax-1] << endl;;
    }
    *ofile << "        </DataArray>" << endl;
    *ofile << "        <CellData>" << endl;
    *ofile << "          <DataArray type=\"Float64\" Name=\"u\" NumberOfTuples=\"" << nVertices << "\" format=\"ascii\" RangeMin=\"" << varMin << "\" RangeMax=\"" << varMax << "\">" << endl;
    for (int i = 0; i < nVertices; i += 6)
    {
        *ofile << "          ";
        int jmax = (i + 6 < nVertices) ? i + 6 : nVertices;
        for (int j = i; j < jmax - 1; j++)
            *ofile << var[j] << " ";
        *ofile << var[jmax-1] << endl;;
    }
    *ofile << "          </DataArray>" << endl;
    *ofile << "        </CellData>" << endl;
    *ofile << "      </Tree>" << endl;
    *ofile << "    </Trees>" << endl;
    *ofile << "  </HyperTreeGrid>" << endl;
    *ofile << "</VTKFile>" << endl;
}

// ****************************************************************************
//  Method: HTGWrite
//
//  Purpose:
//      Creates a HyperTreeGrid and writes it to a file.
//
//  Programmer: Eric Brugger
//  Creation:   Mon May 1 15:28:30 PST 2023
//
// ****************************************************************************

void HTGWrite(const std::string &path, float blankValue, int nx,
    const double *bounds, const float *value)
{
    //
    // Determine the number of levels.
    //
    int nLevels = 1;
    int dim = nx;
    while (dim > 1)
    {
        nLevels++;
        dim = dim / 2;
    }

    //
    // Calculate the total number of vertices.
    //
    int nVertices = 0;
    for (int i = 0; i < nLevels; i++)
        nVertices += 1 << i * 3;

    //
    // Calculate min and max for the variable. We only need to do
    // the input array, since the output will contain the input and
    // averages of the input. We exclude any blank values. If the
    // array contains all blank values we throw an exception.
    //
    int nvals = nx * nx * nx;

    //
    // Find the first non blank value.
    //
    int iReal;
    for (iReal = 0; iReal < nvals && value[iReal] == blankValue; iReal++)
        /* Do nothing */;

    if (iReal == nvals)
    {
        EXCEPTION1(ImproperUseException,
            "HTG Export: The variable only had blank values.");
    }

    float varMin = value[iReal];
    float varMax = value[iReal];
    for (int i = iReal + 1; i < nvals; i++)
    {
        if (value[i] != blankValue)
        {
            varMin = value[i] < varMin ? value[i] : varMin;
            varMax = value[i] > varMax ? value[i] : varMax;
        }
    }

    //
    // Set the number of vertices in each level.
    //
    int *nbVerticesByLevel = new int[nLevels];
    for (int i = 0; i < nLevels; i++)
        nbVerticesByLevel[i] = 1 << i * 3;

    //
    // Create the HTG, specifically the output variable and the mask.
    //
    int *mask = new int[nVertices];
    float *var = new float[nVertices];

    int iLevel = 1;
    int *offsets = new int[nLevels];
    offsets[0] = 0;
    for (int i = 1; i < nLevels; i++)
       offsets[i] = offsets[i-1] + nbVerticesByLevel[i-1];
    int iStart = 0;
    int jStart = 0;
    int kStart = 0;
    var[0] = HTGCreate(value, var, mask, nLevels, nx, blankValue, iLevel,
        offsets, iStart, jStart, kStart);
    mask[0] = var[0] == blankValue ? 1 : 0;
    delete [] offsets;

    //
    // Compress the output variable based on the mask variable.
    //
    int *mask2 = new int[nVertices];
    int nVertices2 = 9;
    int iVar = 9;
    int iMask = 1;
    for (int i = 0; i < 9; i++)
        mask2[i] = mask[i];
    for (int i = 1; i < nLevels - 1; i++)
    {
        int nBits = 1 << i * 3;
        int nbVertices = 0;
        for (int j = 0; j < nBits; j++)
        {
             if (mask[iMask] == 0)
             {
                 nbVertices += 8;
                 for (int k = 0; k < 8; k++)
                 {
                     mask2[nVertices2] = mask[iVar];
                     var[nVertices2] = var[iVar];
                     nVertices2++;
                     iVar++;
                 }
             }
             else
             {
                 iVar += 8;
             }
             iMask++;
        }
        nbVerticesByLevel[i+1] = nbVertices;
    }
    nVertices = nVertices2;
    delete [] mask;
    mask = mask2;

    int nbVerticesByLevelMax = nbVerticesByLevel[nLevels-1];

    //
    // Determine the size of the mask variable. Remove any trailing zeros.
    //
    int lastZero = -1;
    int lastOne = -1;
    for (int i = 0; i < nVertices; i++)
    {
        if (mask[i] == 0)
            lastZero = i;
        else
            lastOne = i;
    }
    int nMask = 1;
    int maskMin = 0;
    int maskMax = 0;
    if (lastOne != -1)
    {
        maskMax = 1;
        nMask = lastOne + 1;
    }

    //
    // Create the descriptor variable. It is the opposite of the mask
    // variable for all but the last level.
    //
    int nDescriptor = 0;
    for (int i = 0; i < nLevels-1; i++)
        nDescriptor += nbVerticesByLevel[i];

    int *descriptor = new int[nDescriptor];
    for (int i = 0; i < nDescriptor; i++)
        descriptor[i] = (mask[i] == 0) ? 1 : 0;

    //
    // Determine the size of the descriptor variable. Remove any trailing zeros.
    //
    lastZero = -1;
    lastOne = -1;
    for (int i = 0; i < nDescriptor; i++)
    {
        if (descriptor[i] == 0)
            lastZero = i;
        else
            lastOne = i;
    }
    nDescriptor = 1;
    int descriptorMin = 0;
    int descriptorMax = 0;
    if (lastZero == -1)
        descriptorMin = 1;
    if (lastOne != -1)
    {
        descriptorMax = 1;
        nDescriptor = lastOne + 1;
    }

    //
    // Replace any blankValue with zero.
    //
    for (int i = 0; i <nVertices; i++)
    {
        if (var[i] == blankValue)
            var[i] = 0.;
    }

    HTGWriteFile(path, bounds, nLevels, nVertices, nDescriptor,
                 descriptorMin, descriptorMax, descriptor,
                 nbVerticesByLevelMax, nbVerticesByLevel,
                 nMask, maskMin, maskMax, mask, varMin, varMax, var);
}
